'''
Created on Aug 23, 2009

@author: mkiyer
'''

from coverage import gen_fw_read_track_from_alignments
from veggie.io.cbowtieparser import BowtieParser
from veggie.sample.sampledb2 import get_sampledb
from veggie.sample.samplegroup import parse_samplegroups_xml
from optparse import OptionParser
import glob
import h5py
import logging
import numpy as np
import os
import sys

def bowtie_to_coverage(analysis_id, bowtie_files, path_to_hdf5,
                       unique_only=True, **kwargs):
    # build read track for alignments in input file
    track = None
    for bowtie_file in bowtie_files:
        reader = BowtieParser(open(bowtie_file))
        logging.debug("        parsing bowtie alignment file: %s" % bowtie_file)
        track = gen_fw_read_track_from_alignments(reader, track=track)            
    # reopen the hdf5 file
    h5file = h5py.File(path_to_hdf5)
    # create hdf5 subgroup for this sample name
    logging.debug("        creating analysis group %s" % analysis_id)
    # create subgroup for this analysis
    analysis_group = h5file.create_group(analysis_id)
    # pile-up reads to compute coverage and write to hdf5 file
    for chrom in track.get_chrom_names():
        logging.debug("        computing coverage on chrom %s" % chrom)
        track.gen_coverage_range_hdf5(analysis_group, chrom,
                                      unique_only=unique_only,
                                      read_length_norm=True,
                                      reads_per_million=True)
    # store attributes for this analysis
    logging.debug("        adding attributes")
    analysis_group.attrs['read_length'] = track.read_length
    read_total = track.total_unique if unique_only else track.total
    analysis_group.attrs['read_total'] = read_total       
    analysis_group.attrs['unique_only'] = unique_only
    analysis_group.attrs['unique_reads'] = track.total_unique
    analysis_group.attrs['total_reads'] = track.total
    
    for k,v in analysis_group.attrs.iteritems():
        logging.debug('            %s: %s' % (k, v))
    # finished this analysis
    del track
    h5file.close()

def build_coverage_tracks(fc_lanes, bowtie_path, path_to_hdf5,
                          **kwargs):
    for flowcell, lane in fc_lanes:
        # make unique name for each analysis
        analysis_id = '%s_%d' % (flowcell, lane)
        logging.debug("analysis: %s" % (analysis_id))
        # open and close the file for each sample so that the hdf5 file can
        # be used by other programs while still being written to
        h5file = h5py.File(path_to_hdf5)        
        # check if this sample has already been processed
        if analysis_id in h5file:
            logging.debug("    hdf5 coverage already exists for analysis %s, skipping.." % analysis_id)
            continue
        # close file immediately so that program can be killed arbitrarily without corrupting
        # the hdf5 file
        h5file.close()
        # search for bowtie results for this analysis
        # look for paired-end results
        fc_lane_pair_glob = '%s_*_single-best.txt' % (analysis_id)
        paired_end_files = glob.glob(os.path.join(bowtie_path, fc_lane_pair_glob))
        if len(paired_end_files) > 0:
            logging.debug("        found paired-end files %s" % ','.join(paired_end_files))
        # look for single read results
        fc_lane_glob = '%s_single-best.txt' % (analysis_id)
        single_read_files = glob.glob(os.path.join(bowtie_path, fc_lane_glob))
        if len(single_read_files) > 0:
            logging.debug("        found single read files %s" % ','.join(single_read_files))
        #fc_lane_glob = '%s_%d*_single-best.txt' % (flowcell, lane)
        # TODO: user should be able to decide whether they want paired-end or single read in
        # situations where both exist.  For now prioritize paired-end vs single-read
        bowtie_files = []
        if len(paired_end_files) > 0:
            bowtie_files.extend(paired_end_files)
        elif len(single_read_files) > 0:
            bowtie_files.extend(single_read_files)
        else:
            logging.debug("    error: no bowtie files for analysis %s, skipping..." % (analysis_id))
            continue
        logging.debug("        using files %s" % ','.join(bowtie_files))
        # build coverage map from bowtie files
        bowtie_to_coverage(analysis_id, bowtie_files, path_to_hdf5,
                           **kwargs)


def get_all_analyses():
    # query sampleDB for all flowcell/lanes
    fclanes = get_sampledb().get_analyses({'app_type': 'Transcriptome'})
    return fclanes

def get_analyses_from_samples(sample_xmlfile):
    fclanes = []
    for sample_group in parse_samplegroups_xml(options.samples):
        for sample in sample_group.samples:
            sample_fclanes = get_sampledb().params[sample]['lanes']
            logging.debug('    sample: %s lanes: %s' %
                          (sample, sample_fclanes))
            fclanes.extend(sample_fclanes)
    return fclanes

if __name__ == '__main__':
    #logging.basicConfig(filename='buildcoveragetracks.log', level=logging.DEBUG)
    logging.basicConfig(level=logging.DEBUG)
    
    optionparser = OptionParser("usage: %prog [options] [-s <samples.xml>]")
    optionparser.add_option('--bowtie-path', dest='bowtie_path',
                            default='/archive10/bigwig/alignments/bowtie-0.10.1/single-best',
                            help="path (absolute or relative) to the directory"
                                 " containing bowtie alignment results"
                                 " [default=%default]")
    optionparser.add_option("-s", "--samples", dest="samples",
                            help="process only the samples defined in the sample XML file")                            
    optionparser.add_option("-o", "--output", dest="outfile",
                            help="output file in HDF5 format")
    optionparser.add_option("-d", "--duplicates", action="store_false", 
                            dest="unique_only", default=True,
                            help="allow duplicate reads")
                            
    (options, args) = optionparser.parse_args()
    if options.samples is not None:
        fclanes = get_analyses_from_samples(options.samples)
    else:
        fclanes = get_all_analyses()
        
    logging.debug('bowtie_path: %s', options.bowtie_path)
    logging.debug('samples: %s', options.samples)
    logging.debug('output_file: %s', options.outfile)
    logging.debug('unique_only: %s', options.unique_only)
    
    build_coverage_tracks(fclanes, options.bowtie_path, options.outfile,
                          unique_only=options.unique_only)
